# -*- coding: utf-8 -*-



import mne
import os,pathlib
import sys
import numpy as np
import scipy.io as sio
import random
import matplotlib.pyplot as plt
import pywt
import gc

# 提取epoch数据
# 提取epoch数据
def parseBci42aFile(dataPath, labelPath, epochWindow=[2,6-1/250]):
    raw_gdf = mne.io.read_raw_gdf(dataPath, stim_channel=['EOG-left', 'EOG-central', 'EOG-right'],preload=True)
    raw_gdf.rename_channels({'EEG-Fz': 'Fz', 'EEG-0': 'FC3', 'EEG-1': 'FC1', 'EEG-2': 'FCz', 'EEG-3': 'FC2', 'EEG-4': 'FC4',
                'EEG-5': 'C5', 'EEG-C3': 'C3', 'EEG-6': 'C1', 'EEG-Cz': 'Cz', 'EEG-7': 'C2', 'EEG-C4': 'C4', 'EEG-8': 'C6',
                'EEG-9': 'CP3', 'EEG-10': 'CP1', 'EEG-11': 'CPz', 'EEG-12': 'CP2', 'EEG-13': 'CP4',
                'EEG-14': 'P1', 'EEG-15': 'Pz', 'EEG-16': 'P2', 'EEG-Pz': 'POz'})
    montage=mne.channels.make_standard_montage("standard_1020")
    raw_gdf.set_montage(montage)
    raw_gdf.drop_channels(['EOG-left', 'EOG-central', 'EOG-right']) # 删除无用的导联
    channel_names = raw_gdf.ch_names
    # raw_resampled = raw_gdf.resample(sfreq=128)
    # 进行[4-38]的带通滤波
    raw_gdf.filter(l_freq=4, h_freq=40,method='iir', fir_window='blackman',iir_params=dict(order=10, ftype='butter'))
    # raw_gdf.plot_sensors(ch_type='eeg', show_names=True)
    # raw_gdf.set_montage(montage)
    # ica去伪迹
    ica = mne.preprocessing.ICA(n_components=22, method='infomax', max_iter='auto')
    ica.fit(raw_gdf)
    # ica.plot_sources(data)
    # ica.plot_properties(data, picks=list(range(0, 15, 1)))
    # 找到伪迹成分
    muscle_idx_auto, scores = ica.find_bads_muscle(raw_gdf)
    eog_idx_auto, eog_scores = ica.find_bads_eog(raw_gdf, ch_name=raw_gdf.ch_names)
    # ecg_idx_auto, ecg_scores = ica.find_bads_ecg(raw_gdf,ch_name=raw_gdf.ch_names)
    # # 去除伪迹
    eog_idx_auto.extend(muscle_idx_auto)
    # ica.exclude = muscle_idx_auto
    ica.exclude = eog_idx_auto
    # ica.exclude = ecg_idx_auto
    ica.apply(raw_gdf)
    # 提取任务标注点
    # mne.events_from_annotations(raw_gdf)返回一个元组，元组的第一个元素保存一个array格式的标注信息，第二个元素包含一个字典标注信息，对应文章标签
    # {'1023': 1, '1072': 2, '276': 3, '277': 4, '32766': 5, '768': 6, '769': 7, '770': 8, '771': 9, '772': 10}
    # 提取任务点数据
    events, event_id = mne.events_from_annotations(raw_gdf)
    epochs = mne.Epochs(raw_gdf, events, event_id=event_id['768'], tmin=epochWindow[0], tmax=epochWindow[1], baseline=(epochWindow[0],epochWindow[0]+0.2),
                preload=True)
    epochs_data = epochs.get_data()
    # # 时频呈现
    # freqs = np.linspace(4, 30, 27)
    # n_cycles = freqs/2.0
    # power_data = mne.time_frequency.tfr_array_morlet(epochs, sfreq=epochs.info['sfreq'],freqs=freqs, n_cycles=n_cycles,use_fft=True, output='power')
    labels = sio.loadmat(labelPath)['classlabel'].squeeze()
    # change the labels from [1-4] to [0-3]
    labels = labels - 1
    data = epochs_data * 1e6
    fs = raw_gdf.info['sfreq']
    return {'data':data, 'labels':labels, 'channel_names':channel_names, 'fs':fs}



def parseBci42aDataset(datasetPath, savePath, epochWindow=[2.5,5.5], verbos=False, isTrans = True):
    '''
    Parse the BCI competition IV-2a data in a MATLAB formate that will be used in the next analysis
    -------
    :param datasetPath: str
        Path to the BCI IV2a original dataset in gdf format.
    :param savePath: str
        Path on where to save the epoched eeg data in a mat format.
    :param epochWindow: list, optional
        time segment to extract in seconds. The default is [0,4].
    :param chans: list
        channels to select from the data.
    :return: None.
    -------
    The dataset will be saved at savePath.
    '''
    subjects = ['A01T', 'A02T', 'A03T', 'A04T', 'A05T', 'A06T', 'A07T', 'A08T', 'A09T', 'A01E', 'A02E', 'A03E', 'A04E',
                'A05E', 'A06E', 'A07E', 'A08E', 'A09E']  # 训练集

    print('Extracting the data into mat format:')
    if not os.path.exists(savePath):
        os.makedirs(savePath)
    print('Processed data be saved in folder:' + savePath)
    for iSubs, subs in enumerate(subjects):  # s, se
        if not os.path.exists(os.path.join(datasetPath, subs + '.mat')):
            raise ValueError('The BCI-IV-2a original dataset doesn\'t exist at path:' +
                             os.path.join(datasetPath, subs + '.mat') +
                             'Please download and copy the extracted dataset at the above path.\n' +
                             'More details about how to download this data can be found in the Instructions.txt file')

        print('Processing subject No:' + subjects[iSubs])
        data = parseBci42aFile(os.path.join(datasetPath, subs + '.gdf'),
                               os.path.join(datasetPath, subs + '.mat'),
                               epochWindow=epochWindow)

        sio.savemat(os.path.join(savePath, subjects[iSubs] + '.mat'), data)

parseBci42aDataset(r'D:\Project_mb\data\BCI42a\raw_data',r'D:\Project_mb\data\BCI42a\preprocession_matdata',epochWindow=[1.5,6])
